[mlir][xegpu] Add definition of SliceAttr #150146

chencha3 · 2025-07-22T23:47:16Z

This PR introduces the definition of SliceAttr, which complements the existing LayoutAttr by providing a mechanism to describe data distribution for operations such as multi_reduction and broadcast. For potential convenience of lowering pass, A common interface (LayoutTrait) is created for LayoutAttr and SliceAttr to provide an util to compute offsets for a given subgroup (id), and workgroup-level problem size (shape). the WgToSgCreateNdOp is updated to test the implementation of such interface for LayoutAttr, and a test pass --test-xegpu-layout-interface is introduced to test its implementation for SliceAttr.
Additional use cases and lowering support for SliceAttr will be addressed in subsequent PRs.

github-actions · 2025-07-23T20:36:19Z

✅ With the latest revision this PR passed the C/C++ code formatter.

Jianhui-Li · 2025-07-25T20:40:12Z

mlir/include/mlir/Dialect/XeGPU/IR/XeGPUDialect.td

+      /// drops the data in the specified dimension, and return the rest. e.g.,
+      /// for data = [32, 64, 8], dropPositions = [0, 2], it will return [64]
+      template<typename T, typename U>
+      static llvm::SmallVector<T> dropDims(llvm::ArrayRef<T> data, llvm::ArrayRef<U> dropPositions) {


nit : dropDims -> sliceDims? data->dims, dropPosition -> dropDims

changed it to static llvm::SmallVector<T> slice(llvm::ArrayRef<T> shape, llvm::ArrayRef<U> dims). is it good to you?

Jianhui-Li · 2025-07-25T21:13:45Z

mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp

+        Value add = rewriter.createOrFold<index::AddOp>(
+            loc, tdescOffsets[i],
+            getValueOrCreateConstantIndexOp(rewriter, loc, oldOffsets[idx]));
+        newOffsets.push_back(add);


nit: oldOffsets: wgTileOffsets
newOffsets: sgTileOffsets

good suggestion. Fixed.

Jianhui-Li · 2025-07-25T21:33:54Z

mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td

  }];

  let assemblyFormat = "`<` struct(params) `>`";
  let genVerifyDecl = 1;
 }

+
+def XeGPU_SliceAttr : XeGPUAttr<"Slice", "slice", [LayoutTrait]> {


I think it is desirable to allow nested slice attribute to match the staged reduction use case, where reduction may follow another reduction.

will add it.

the support for nested SliceAttr has been enabled. @Jianhui-Li

nbpatel · 2025-08-04T15:24:52Z

mlir/test/Dialect/XeGPU/layout.mlir

@@ -50,4 +50,21 @@ gpu.func @convert_layout_wg(%a: vector<32x64xf16>) {
  gpu.return
 }

+gpu.func @slice_attr_repeat_dim() {
+  //CHECK: arith.constant {layout_result_0 = #xegpu.slice<<sg_layout = [16, 1, 1], sg_data = [1, 8, 2]>, dims = [2]>} dense<8> : vector<16x8xindex>
+  %cst = arith.constant {layout_result_0 = #xegpu.slice<<sg_layout = [16, 1, 1], sg_data = [1, 8, 2]>, dims = [2]>} dense<8> : vector<16x8xindex>


shouldn't it be #xegpu.slice<<#xegpu.layout<sg_layout = [16, 1, 1], sg_data = [1, 8, 2]>>, dims = [2]>}

the print format issue. it is fixed.

nbpatel · 2025-08-04T15:28:01Z

mlir/lib/Dialect/XeGPU/IR/CMakeLists.txt

@@ -7,6 +7,7 @@ add_mlir_dialect_library(MLIRXeGPUDialect

  DEPENDS
  MLIRXeGPUIncGen
+  MLIRXeGPUAttrInterfaceIncGen
  MLIRXeGPUAttrsIncGen
  MLIRXeGPUEnumsIncGen



dont you need to link MLIRIndexDialect ,MLIRAffineUtils here?

Good catch. Yes, they are needed.

charithaintc

LGTM % comments.

charithaintc · 2025-08-05T18:55:56Z

mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td

+                    "delinearizeSubgroupId",
+                    (ins "OpBuilder &": $builder, "Location":$loc, "Value":$linearId)>,
+    InterfaceMethod<"Get the local offset to be accessed by the given subgroup Id",
+                    "FailureOr<SmallVector<SmallVector<Value>>>",


what is the need for having a vector<vector<>> here?

Since each subgroup may be assigned with multiple blocks.

charithaintc · 2025-08-05T19:16:12Z

mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td

+    #r = #xegpu.slice<#l, dim = 0>
+
+    %exp = math.exp %input {layout_result_0 = #l}: vector<256x128xf32>
+    %red = vector.multi_reduction<add>, %exp, %acc [0] {layout_result_0 = #r}: vector<256x128xf32> to vector<128xf32>


better to add a comment here explaning the output layout of %red.

here the output sg layout is still [8, 4]. however data is shared along dim 0. So effectively there are 8 slices of [1, 4] SG segments each owning [1, 32] data. For example SGs [0-7][0] owns the same 1x32 data segement.

mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td

charithaintc · 2025-08-05T19:22:14Z

mlir/include/mlir/Dialect/XeGPU/IR/XeGPUDialect.td

+      /// drops/slices the shape in the specified dims, and return the rest. e.g.,
+      /// for shape = [32, 64, 8], dims = [0, 2], it will return [64]
+      template<typename T, typename U>
+      static llvm::SmallVector<T> slice(llvm::ArrayRef<T> shape, llvm::ArrayRef<U> dims) {


seems like this unrelated to XeGPUDialect. any reason for placing here? can it be moved to Utils?

Moving it to Utils will incur recursive dependence for linker.

charithaintc · 2025-08-05T19:25:59Z

mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp

@@ -33,6 +37,54 @@ void XeGPUDialect::initialize() {
      >();
 }

+// generate offsets computing instructions for a subgroup


nit: Use triple line comments for def/impl.

charithaintc · 2025-08-05T19:33:49Z

mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp

+  if (!isWgLayout())
+    return failure();
+
+  // TODO: handle order attribute


maybe assert unavailability of Order attribute?

charithaintc · 2025-08-05T19:35:04Z

mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp

+}
+
+FailureOr<SmallVector<SmallVector<Value>>>
+LayoutAttr::getOffsets(OpBuilder &builder, Location loc, Value linearId,


Better to add a comment on what the purpose of this function.

charithaintc · 2025-08-05T19:42:06Z

mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp

+  // check every element in dims is unique and smaller than rank
+  llvm::SmallDenseSet<int64_t> seen;
+  for (int64_t dim : dims.asArrayRef()) {
+    if (dim >= rank)


nit: should we check if dim >= 0 here?

good catch. Fixed.

charithaintc · 2025-08-05T19:52:35Z

mlir/test/lib/Dialect/XeGPU/TestXeGPUTransforms.cpp

+#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ")
+#define LDBG(X) LLVM_DEBUG(DBGS() << X << "\n")
+
+class TestStepOpPattern : public OpConversionPattern<vector::StepOp> {


please add a comment on what this pattern is doing.

Yes, fixed.

Co-authored-by: Charitha Saumya <[email protected]>

nbpatel

LGTM

nbpatel · 2025-08-06T17:47:43Z

mlir/include/mlir/Dialect/XeGPU/IR/CMakeLists.txt

+mlir_tablegen(XeGPUAttrInterface.h.inc -gen-attr-interface-decls)
+mlir_tablegen(XeGPUAttrInterface.cpp.inc -gen-attr-interface-defs)
+add_public_tablegen_target(MLIRXeGPUAttrInterfaceIncGen)
+add_dependencies(mlir-headers MLIRXeGPUAttrInterfaceIncGen)


nbpatel · 2025-08-06T17:54:35Z

mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp

-
-      auto newCreateNdOp = xegpu::CreateNdDescOp::create(
-          rewriter, loc, newTdescTy, op.getSource(), globalOffsets,
+    SmallVector<OpFoldResult> wgTileOffsets = op.getMixedOffsets();


nit: maybe better to not use tile and just use wgOffsets and sgOffsets as var names

nbpatel · 2025-08-06T17:55:29Z

mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td

+  }];
+
+  let parameters = (ins
+    "xegpu::LayoutTrait": $parent,


why is parent a trait and not LayoutAttr?

It is to support using both LayoutAttr and SliceAttr as parent. The later one is nested definition.

Jianhui-Li · 2025-08-07T22:09:56Z

mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp

+static SmallVector<SmallVector<Value>>
+genOffsetsComputations(OpBuilder &builder, Location loc,
+                       SmallVector<Value> sgId, ArrayRef<int64_t> sgLayout,
+                       ArrayRef<int64_t> sgShape, ArrayRef<int64_t> shape) {


consider generalizing sgShape to sgDataStride.
From the caller size, need to set as following :
sgDataStride[i] = sgData[i] if (wgData[i] % (sgData[i]*sgLayout[i])) == 0
sgDataStride[i] = 0 if wgData[i] == sgData[i]

Jianhui-Li · 2025-08-07T22:23:31Z

mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp

+  if (auto maybeSgShape = getSgDataAsInt())
+    sgShape = maybeSgShape.value();
+  else if (auto ratio = computeShapeRatio(shape, sgLayout))
+    sgShape = ratio.value();


shape == ratio?

changed the variable name to derivedShape for clarification per discussion.

It seems you change the last place, but missed this one.

Jianhui-Li · 2025-08-07T22:35:36Z

mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir

+    //CHECK: [[C128:%.+]] = arith.constant 128 : index
+    //CHECK: [[X:%.+]] = index.remu [[UX]], [[C128]]
+    //CHECK: [[TDESC:%.+]] = xegpu.create_nd_tdesc [[ARG_0]][[[Y]], [[X]]] : memref<256x128xf32> -> !xegpu.tensor_desc<32x32xf32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
+    %tdesc = xegpu.create_nd_tdesc %src[0, 0] : memref<256x128xf32>


consider adding a test:

%tdesc = xegpu.create_nd_tdesc %src[0, 0] : memref<256x32xf32>
-> !xegpu.tensor_desc<256x32xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [1, 16], lane_data = [1, 1]>>

Yes, added one.

Jianhui-Li

LGTM

#1104) Brings Slice attribute def (llvm/llvm-project#150146)

chencha3 added 10 commits July 22, 2025 19:46

add definition draft of SliceAttr

2bc70b6

add layout traits

3959f9e

add verifier and interface

2027cfc

add invalid unit test

638c085

add wrappers

91048f0

Merge branch 'main' into xegpu_slice_attr

7eaf0a6

update description

ddc42c2

refactor

36e2c3a

add delinearizeSubgroupId interface

6872e6d

Merge branch 'main' into xegpu_slice_attr

ded53b4

chencha3 added 8 commits July 23, 2025 21:05

fix format

223fab9

add impl of getOffsets for LayoutAttr

60e20a0

apply getOffsets in CreateNdDescOp

3630966

cleanup

398d69b

fix a bug

08e4aa9

cleanup

62aa1dd

add unit test

de0a1bb

Merge branch 'main' into xegpu_slice_attr

a483699

chencha3 marked this pull request as ready for review July 25, 2025 17:28

fix a typo

e7f2977

Jianhui-Li reviewed Jul 25, 2025

View reviewed changes

add unit test

e3e4a61

nbpatel reviewed Aug 4, 2025

View reviewed changes

chencha3 added 4 commits August 4, 2025 16:03

Merge branch 'main' into xegpu_slice_attr

4d72663

fix conflicts

3f59105

address comments

129312a

add support for nested SliceAttr

0865612

chencha3 added 2 commits August 5, 2025 17:30

add unit test for nested slice attr

b67f2b1

cleanup

01e4efe

charithaintc self-requested a review August 5, 2025 18:38

charithaintc approved these changes Aug 5, 2025

View reviewed changes

chencha3 and others added 3 commits August 5, 2025 16:37

Update mlir/include/mlir/Dialect/XeGPU/IR/XeGPUAttrs.td

3077c6c

Co-authored-by: Charitha Saumya <[email protected]>

update docs

d1f7bac

add check for order attribute

27da02a

chencha3 requested review from Jianhui-Li and nbpatel August 6, 2025 17:46

clean up

e49e1cf

nbpatel approved these changes Aug 6, 2025

View reviewed changes

clean up

59de450

Jianhui-Li reviewed Aug 7, 2025

View reviewed changes

chencha3 added 2 commits August 8, 2025 15:07

address comments

1b16552

cleanup

0511e1b

Jianhui-Li approved these changes Aug 8, 2025

View reviewed changes

chencha3 merged commit c962234 into llvm:main Aug 8, 2025
9 checks passed

rupprecht added a commit to rupprecht/llvm-project that referenced this pull request Aug 8, 2025

[bazel] Port llvm#150146: xegpu/sliceattr

3e4dbd2

rupprecht added a commit that referenced this pull request Aug 8, 2025

[bazel] Port #150146: xegpu/sliceattr (#152772)

4e58724

Garra1980 mentioned this pull request Aug 13, 2025

[LLVM Pulldown] Bump to LLVM rev 92164faf17d553359418b9f49c1a41d680d0… intel/mlir-extensions#1104

Merged

6 tasks

Garra1980 added a commit to intel/mlir-extensions that referenced this pull request Aug 13, 2025

[LLVM Pulldown] Bump to LLVM rev 92164faf17d553359418b9f49c1a41d680d0… (

70e160a

#1104) Brings Slice attribute def (llvm/llvm-project#150146)

[mlir][xegpu] Add definition of SliceAttr #150146

[mlir][xegpu] Add definition of SliceAttr #150146

Uh oh!

Conversation

chencha3 commented Jul 22, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions bot commented Jul 23, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Jianhui-Li Jul 25, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

nbpatel Aug 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

charithaintc left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

nbpatel left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

chencha3 commented Jul 22, 2025 •

edited

Loading

github-actions bot commented Jul 23, 2025 •

edited

Loading

Jianhui-Li Jul 25, 2025 •

edited

Loading

nbpatel Aug 4, 2025 •

edited

Loading

Jianhui-Li Aug 7, 2025 •

edited by chencha3

Loading